{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "scratchpad",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/margaretmz/CartoonGAN-e2e-tflite-tutorial/blob/master/ml/metadata/Add_Metadata.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "laCUI2hzapKv",
        "colab_type": "text"
      },
      "source": [
        "- Reference: https://github.com/margaretmz/selfie2anime-with-tflite/blob/master/ml/add-meta-data-Colab/Add%20metadata%20to%20selfie2anime.ipynb. \n",
        "- TensorFlow Lite meatdata: https://www.tensorflow.org/lite/convert/metadata.\n",
        "- Authored by: Sayak and Margaret\n",
        "- Updated on: August 04 2020"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "lIYdn1woOS1n",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        },
        "outputId": "12660303-8384-4ff1-eb5a-2eba967cd2f6"
      },
      "source": [
        "!pip install -q tflite-support"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\u001b[?25l\r\u001b[K     |▉                               | 10kB 17.4MB/s eta 0:00:01\r\u001b[K     |█▊                              | 20kB 1.8MB/s eta 0:00:01\r\u001b[K     |██▌                             | 30kB 2.3MB/s eta 0:00:01\r\u001b[K     |███▍                            | 40kB 2.6MB/s eta 0:00:01\r\u001b[K     |████▏                           | 51kB 2.0MB/s eta 0:00:01\r\u001b[K     |█████                           | 61kB 2.3MB/s eta 0:00:01\r\u001b[K     |█████▉                          | 71kB 2.5MB/s eta 0:00:01\r\u001b[K     |██████▊                         | 81kB 2.7MB/s eta 0:00:01\r\u001b[K     |███████▋                        | 92kB 2.8MB/s eta 0:00:01\r\u001b[K     |████████▍                       | 102kB 2.8MB/s eta 0:00:01\r\u001b[K     |█████████▎                      | 112kB 2.8MB/s eta 0:00:01\r\u001b[K     |██████████                      | 122kB 2.8MB/s eta 0:00:01\r\u001b[K     |███████████                     | 133kB 2.8MB/s eta 0:00:01\r\u001b[K     |███████████▊                    | 143kB 2.8MB/s eta 0:00:01\r\u001b[K     |████████████▋                   | 153kB 2.8MB/s eta 0:00:01\r\u001b[K     |█████████████▍                  | 163kB 2.8MB/s eta 0:00:01\r\u001b[K     |██████████████▎                 | 174kB 2.8MB/s eta 0:00:01\r\u001b[K     |███████████████▏                | 184kB 2.8MB/s eta 0:00:01\r\u001b[K     |████████████████                | 194kB 2.8MB/s eta 0:00:01\r\u001b[K     |████████████████▉               | 204kB 2.8MB/s eta 0:00:01\r\u001b[K     |█████████████████▋              | 215kB 2.8MB/s eta 0:00:01\r\u001b[K     |██████████████████▌             | 225kB 2.8MB/s eta 0:00:01\r\u001b[K     |███████████████████▎            | 235kB 2.8MB/s eta 0:00:01\r\u001b[K     |████████████████████▏           | 245kB 2.8MB/s eta 0:00:01\r\u001b[K     |█████████████████████           | 256kB 2.8MB/s eta 0:00:01\r\u001b[K     |█████████████████████▉          | 266kB 2.8MB/s eta 0:00:01\r\u001b[K     |██████████████████████▊         | 276kB 2.8MB/s eta 0:00:01\r\u001b[K     |███████████████████████▌        | 286kB 2.8MB/s eta 0:00:01\r\u001b[K     |████████████████████████▍       | 296kB 2.8MB/s eta 0:00:01\r\u001b[K     |█████████████████████████▏      | 307kB 2.8MB/s eta 0:00:01\r\u001b[K     |██████████████████████████      | 317kB 2.8MB/s eta 0:00:01\r\u001b[K     |██████████████████████████▉     | 327kB 2.8MB/s eta 0:00:01\r\u001b[K     |███████████████████████████▊    | 337kB 2.8MB/s eta 0:00:01\r\u001b[K     |████████████████████████████▌   | 348kB 2.8MB/s eta 0:00:01\r\u001b[K     |█████████████████████████████▍  | 358kB 2.8MB/s eta 0:00:01\r\u001b[K     |██████████████████████████████▎ | 368kB 2.8MB/s eta 0:00:01\r\u001b[K     |███████████████████████████████ | 378kB 2.8MB/s eta 0:00:01\r\u001b[K     |████████████████████████████████| 389kB 2.8MB/s eta 0:00:01\r\u001b[K     |████████████████████████████████| 399kB 2.8MB/s \n",
            "\u001b[?25h  Building wheel for tflite-support (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hSBlosV-Wim7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "import tensorflow as tf\n",
        "from absl import flags"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mayO_BGeWrT1",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from tflite_support import flatbuffers\n",
        "from tflite_support import metadata as _metadata\n",
        "from tflite_support import metadata_schema_py_generated as _metadata_fb"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v_czMsARW1SW",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!mkdir model_without_metadata\n",
        "!mkdir model_with_metadata"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8_Qa9lo4W6zV",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 578
        },
        "outputId": "862985f5-0d85-4522-9f8f-e60107981146"
      },
      "source": [
        "!wget https://storage.googleapis.com/cartoon_gan/fixed_shaped_models/whitebox_cartoon_gan_dr.tflite\n",
        "!wget https://storage.googleapis.com/cartoon_gan/fixed_shaped_models/whitebox_cartoon_gan_fp16.tflite\n",
        "!wget https://storage.googleapis.com/cartoon_gan/fixed_shaped_models/whitebox_cartoon_gan_int8.tflite\n",
        "\n",
        "!mv *.tflite model_without_metadata/"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "--2020-08-04 15:41:44--  https://storage.googleapis.com/cartoon_gan/fixed_shaped_models/whitebox_cartoon_gan_dr.tflite\n",
            "Resolving storage.googleapis.com (storage.googleapis.com)... 173.194.213.128, 173.194.216.128, 74.125.26.128, ...\n",
            "Connecting to storage.googleapis.com (storage.googleapis.com)|173.194.213.128|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 2568976 (2.4M) [application/octet-stream]\n",
            "Saving to: ‘whitebox_cartoon_gan_dr.tflite’\n",
            "\n",
            "\r          whitebox_   0%[                    ]       0  --.-KB/s               \rwhitebox_cartoon_ga 100%[===================>]   2.45M  --.-KB/s    in 0.01s   \n",
            "\n",
            "2020-08-04 15:41:44 (187 MB/s) - ‘whitebox_cartoon_gan_dr.tflite’ saved [2568976/2568976]\n",
            "\n",
            "--2020-08-04 15:41:48--  https://storage.googleapis.com/cartoon_gan/fixed_shaped_models/whitebox_cartoon_gan_fp16.tflite\n",
            "Resolving storage.googleapis.com (storage.googleapis.com)... 172.217.204.128, 172.217.203.128, 173.194.216.128, ...\n",
            "Connecting to storage.googleapis.com (storage.googleapis.com)|172.217.204.128|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 3481632 (3.3M) [application/octet-stream]\n",
            "Saving to: ‘whitebox_cartoon_gan_fp16.tflite’\n",
            "\n",
            "whitebox_cartoon_ga 100%[===================>]   3.32M  --.-KB/s    in 0.03s   \n",
            "\n",
            "2020-08-04 15:41:48 (107 MB/s) - ‘whitebox_cartoon_gan_fp16.tflite’ saved [3481632/3481632]\n",
            "\n",
            "--2020-08-04 15:41:53--  https://storage.googleapis.com/cartoon_gan/fixed_shaped_models/whitebox_cartoon_gan_int8.tflite\n",
            "Resolving storage.googleapis.com (storage.googleapis.com)... 172.217.204.128, 172.217.203.128, 173.194.215.128, ...\n",
            "Connecting to storage.googleapis.com (storage.googleapis.com)|172.217.204.128|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 1808784 (1.7M) [application/octet-stream]\n",
            "Saving to: ‘whitebox_cartoon_gan_int8.tflite’\n",
            "\n",
            "whitebox_cartoon_ga 100%[===================>]   1.72M  --.-KB/s    in 0.02s   \n",
            "\n",
            "2020-08-04 15:41:53 (105 MB/s) - ‘whitebox_cartoon_gan_int8.tflite’ saved [1808784/1808784]\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "InH13DM9cdqP",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# This is where we will export a new .tflite model file with metadata, and a .json file with metadata info\n",
        "EXPORT_DIR = \"model_with_metadata\""
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E4ta4ByLYST3",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class MetadataPopulatorForGANModel(object):\n",
        "  \"\"\"Populates the metadata for the CartoonGAN model.\"\"\"\n",
        "\n",
        "  def __init__(self, model_file):\n",
        "    self.model_file = model_file\n",
        "    self.metadata_buf = None\n",
        "\n",
        "  def populate(self):\n",
        "    \"\"\"Creates metadata and then populates it for a style transfer model.\"\"\"\n",
        "    self._create_metadata()\n",
        "    self._populate_metadata()\n",
        "  \n",
        "  def _create_metadata(self):\n",
        "    \"\"\"Creates the metadata for the CartoonGAN model.\"\"\"\n",
        "\n",
        "    # Creates model info.\n",
        "    model_meta = _metadata_fb.ModelMetadataT()\n",
        "    model_meta.name = \"CartoonGAN\" \n",
        "    model_meta.description = (\"Cartoonizes an image. Reference: https://bit.ly/cartoon-gan. TFLiteConverter used from tf-nightly.\")\n",
        "    model_meta.version = \"v1\"\n",
        "    model_meta.author = \"TensorFlow\"\n",
        "    model_meta.license = (\"Apache License. Version 2.0 \"\n",
        "                          \"http://www.apache.org/licenses/LICENSE-2.0.\")\n",
        "\n",
        "    # Creates info for the input, normal image.\n",
        "    input_image_meta = _metadata_fb.TensorMetadataT()\n",
        "    input_image_meta.name = \"source_image\"\n",
        "    # if self.model_type==\"other\":\n",
        "    input_image_meta.description = (\n",
        "            \"The expected image is 512 x 512, with three channels \"\n",
        "            \"(blue, red, and green) per pixel. Each value in the tensor is between\"\n",
        "            \" -1 and 1.\")\n",
        "    # elif self.model_type==\"fp16\":\n",
        "    #     input_image_meta.description = (\n",
        "    #         \"The expected image is 512 x 512, with three channels \"\n",
        "    #         \"(blue, red, and green) per pixel. Each value in the tensor is between\"\n",
        "    #         \" -1 and 1.\")\n",
        "    input_image_meta.content = _metadata_fb.ContentT()\n",
        "    input_image_meta.content.contentProperties = (\n",
        "        _metadata_fb.ImagePropertiesT())\n",
        "    input_image_meta.content.contentProperties.colorSpace = (\n",
        "        _metadata_fb.ColorSpaceType.RGB)\n",
        "    input_image_meta.content.contentPropertiesType = (\n",
        "        _metadata_fb.ContentProperties.ImageProperties)\n",
        "    input_image_normalization = _metadata_fb.ProcessUnitT()\n",
        "    input_image_normalization.optionsType = (\n",
        "        _metadata_fb.ProcessUnitOptions.NormalizationOptions)\n",
        "    input_image_normalization.options = _metadata_fb.NormalizationOptionsT()\n",
        "    input_image_normalization.options.mean = [127.5]\n",
        "    input_image_normalization.options.std = [127.5]\n",
        "    input_image_meta.processUnits = [input_image_normalization]\n",
        "    input_image_stats = _metadata_fb.StatsT()\n",
        "    input_image_stats.max = [1.0]\n",
        "    input_image_stats.min = [-1.0]\n",
        "    input_image_meta.stats = input_image_stats\n",
        "\n",
        "\n",
        "    # Creates output info, cartoonized image\n",
        "    output_image_meta = _metadata_fb.TensorMetadataT()\n",
        "    output_image_meta.name = \"cartoonized_image\"\n",
        "    output_image_meta.description = \"Image cartoonized.\"\n",
        "    output_image_meta.content = _metadata_fb.ContentT()\n",
        "    output_image_meta.content.contentProperties = _metadata_fb.ImagePropertiesT()\n",
        "    output_image_meta.content.contentProperties.colorSpace = (\n",
        "        _metadata_fb.ColorSpaceType.RGB)\n",
        "    output_image_meta.content.contentPropertiesType = (\n",
        "        _metadata_fb.ContentProperties.ImageProperties)\n",
        "    output_image_normalization = _metadata_fb.ProcessUnitT()\n",
        "    output_image_normalization.optionsType = (\n",
        "        _metadata_fb.ProcessUnitOptions.NormalizationOptions)\n",
        "    output_image_normalization.options = _metadata_fb.NormalizationOptionsT()\n",
        "    output_image_normalization.options.mean = [-1]\n",
        "    output_image_normalization.options.std = [0.00784313] # 1/127.5\n",
        "    output_image_meta.processUnits = [output_image_normalization]\n",
        "    output_image_stats = _metadata_fb.StatsT()\n",
        "    output_image_stats.max = [255.0]\n",
        "    output_image_stats.min = [0.0]\n",
        "    output_image_meta.stats = output_image_stats\n",
        "\n",
        "    # Creates subgraph info.\n",
        "    subgraph = _metadata_fb.SubGraphMetadataT()\n",
        "    subgraph.inputTensorMetadata = [input_image_meta] \n",
        "    subgraph.outputTensorMetadata = [output_image_meta] \n",
        "    model_meta.subgraphMetadata = [subgraph]\n",
        "\n",
        "    b = flatbuffers.Builder(0)\n",
        "    b.Finish(\n",
        "        model_meta.Pack(b),\n",
        "        _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)\n",
        "    self.metadata_buf = b.Output()\n",
        "\n",
        "  def _populate_metadata(self):\n",
        "    \"\"\"Populates metadata to the model file.\"\"\"\n",
        "    populator = _metadata.MetadataPopulator.with_model_file(self.model_file)\n",
        "    populator.load_metadata_buffer(self.metadata_buf)\n",
        "    populator.populate()"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vRlUqyFMcNyy",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def populate_metadata(model_file):\n",
        "  \"\"\"Populates the metadata using the populator specified.\n",
        "  Args:\n",
        "      model_file: valid path to the model file.\n",
        "      model_type: a type defined in StyleTransferModelType .\n",
        "  \"\"\"\n",
        "\n",
        "  # Populates metadata for the model.\n",
        "  model_file_basename = os.path.basename(model_file)\n",
        "  export_path = os.path.join(EXPORT_DIR, model_file_basename)\n",
        "  tf.io.gfile.copy(model_file, export_path, overwrite=True)\n",
        "\n",
        "  populator = MetadataPopulatorForGANModel(export_path) \n",
        "  populator.populate()\n",
        "\n",
        "  # Displays the metadata that was just populated into the tflite model.\n",
        "  displayer = _metadata.MetadataDisplayer.with_model_file(export_path)\n",
        "  export_json_file = os.path.join(\n",
        "      EXPORT_DIR,\n",
        "      os.path.splitext(model_file_basename)[0] + \".json\")\n",
        "  json_file = displayer.get_metadata_json()\n",
        "  with open(export_json_file, \"w\") as f:\n",
        "    f.write(json_file)\n",
        "  print(\"Finished populating metadata and associated file to the model:\")\n",
        "  print(export_path)\n",
        "  print(\"The metadata json file has been saved to:\")\n",
        "  print(os.path.join(EXPORT_DIR,\n",
        "                   os.path.splitext(model_file_basename)[0] + \".json\"))"
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YVSJFR4HcUuI",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 85
        },
        "outputId": "30edd015-8a9b-4dda-aa71-9e1d0f887f17"
      },
      "source": [
        "quantization = \"fp16\" #@param [\"dr\", \"int8\", \"fp16\"]\n",
        "tflite_model_path = f\"whitebox_cartoon_gan_{quantization}.tflite\" \n",
        "MODEL_FILE = \"/content/model_without_metadata/{}\".format(tflite_model_path)\n",
        "populate_metadata(MODEL_FILE)"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Finished populating metadata and associated file to the model:\n",
            "model_with_metadata/whitebox_cartoon_gan_fp16.tflite\n",
            "The metadata json file has been saved to:\n",
            "model_with_metadata/whitebox_cartoon_gan_fp16.json\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SClcLn9Pv702",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from google.colab import auth as google_auth\n",
        "google_auth.authenticate_user()"
      ],
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "szJR7LmxwIHw",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 153
        },
        "outputId": "c133bf73-dc97-466e-a81b-5ce946d738bc"
      },
      "source": [
        "!gsutil -m cp -r model_with_metadata/* gs://cartoon_gan/fixed_shaped_models/with_metadata/"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Copying file://model_with_metadata/whitebox_cartoon_gan_dr.json [Content-Type=application/json]...\n",
            "/ [0/6 files][    0.0 B/  7.5 MiB]   0% Done                                    \rCopying file://model_with_metadata/whitebox_cartoon_gan_int8.tflite [Content-Type=application/octet-stream]...\n",
            "/ [0/6 files][    0.0 B/  7.5 MiB]   0% Done                                    \rCopying file://model_with_metadata/whitebox_cartoon_gan_dr.tflite [Content-Type=application/octet-stream]...\n",
            "/ [0/6 files][    0.0 B/  7.5 MiB]   0% Done                                    \rCopying file://model_with_metadata/whitebox_cartoon_gan_fp16.json [Content-Type=application/json]...\n",
            "/ [0/6 files][    0.0 B/  7.5 MiB]   0% Done                                    \rCopying file://model_with_metadata/whitebox_cartoon_gan_fp16.tflite [Content-Type=application/octet-stream]...\n",
            "Copying file://model_with_metadata/whitebox_cartoon_gan_int8.json [Content-Type=application/json]...\n",
            "- [6/6 files][  7.5 MiB/  7.5 MiB] 100% Done                                    \n",
            "Operation completed over 6 objects/7.5 MiB.                                      \n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}